import {
  TCompareFunc,
  defaultCompare,
} from '../utils/tools';
import {
  COLORS_TXT,
  LESS_THAN,
} from '../utils/constant';
import {
  BinarySearchTree,
} from './bstree';
import { RedBlackTreeNode, TRedBlackTreeNode, } from './base';

export class RedBlackTree extends BinarySearchTree {
  public root: TRedBlackTreeNode;

  public constructor(elements: any[], funcCompare: TCompareFunc = defaultCompare) {
    super(elements, funcCompare);
  }

  public insert(key: any) {
    if (!this.root) {
      this.root = new RedBlackTreeNode(key, COLORS_TXT.BLACK);
    } else {
      const newNode = this.insertNode(this.root, key);
      this.fixTreeProperties(newNode);
    }
  }

  public remove(key: any) {
    this.root = this.removeNode(key, this.root) as TRedBlackTreeNode;
    this.fixTreeProperties(this.root);
  }

  protected insertNode(node: RedBlackTreeNode, key: any) {
    if (this.funcCompare(key, node.key) === LESS_THAN) {
      if (!node.left) {
        node.left = new RedBlackTreeNode(key, COLORS_TXT.RED);
        node.left.parent = node;
        return node.left;
      } else {
        this.insertNode(node.left, key);
      }
    } else if (!node.right) {
      node.right = new RedBlackTreeNode(key, COLORS_TXT.RED);
      node.right.parent = node;
      return node.right;
    } else {
      return this.insertNode(node.right, key);
    }
  }

  protected rotateLL(node: RedBlackTreeNode) {
    let currentNode = node;
    const temp = currentNode?.left;
    if (currentNode && temp) {
      currentNode.left = temp.right;

      if (temp.right) {
        temp.right.parent = currentNode;
      }
      temp.parent = currentNode.parent;

      if (!currentNode.parent) {
        this.root = temp;
      } else {
        if (currentNode.parent.left === currentNode) {
          currentNode.parent.left = temp;
        } else if (currentNode.parent.right === currentNode) {
          currentNode.parent.right = temp;
        }
      }

      currentNode.parent = temp;
      temp.right = currentNode;
    }
    return currentNode;
  }

  protected rotateRR(node: RedBlackTreeNode) {
    let currentNode = node;
    const temp = currentNode?.right;
    if (currentNode && temp) {
      currentNode.right = temp.left;

      if (temp.left) {
        temp.left.parent = currentNode;
      }
      temp.parent = currentNode.parent;

      if (!currentNode.parent) {
        this.root = temp;
      } else {
        if (currentNode.parent.left === currentNode) {
          currentNode.parent.left = temp;
        } else if (currentNode.parent.right === currentNode) {
          currentNode.parent.right = temp;
        }
      }

      currentNode.parent = temp;
      temp.left = currentNode;
    }
    return currentNode;
  }

  // protected removeNode(key: any, node: TRedBlackTreeNode) {
  //   let newNode = super.removeNode(key, node) as TRedBlackTreeNode;

  //   this.fixTreeProperties(newNode);
  //   return newNode;
  // }

  protected fixTreeProperties(node: TRedBlackTreeNode) {
    let currentNode = node;
    
    while(currentNode?.isRed() && currentNode?.parent?.isRed()) {
      let parent = currentNode.parent;
      let grandParent = parent.parent;
      
      if (grandParent?.left === parent) {
        const uncle = grandParent.right;
        
        if (uncle?.isRed()) {
          grandParent.color = COLORS_TXT.RED;
          uncle.color = COLORS_TXT.BLACK;
          parent.color = COLORS_TXT.BLACK;
          
          currentNode = grandParent;
        } else {
          if (parent.right === currentNode) {
            currentNode = parent;
            parent = this.rotateRR(parent);
          }
          currentNode = this.rotateLL(grandParent);
          currentNode.color = COLORS_TXT.BLACK;
          grandParent.color = COLORS_TXT.RED;
        }

      } else if (grandParent?.right === parent) {
        const uncle = grandParent.left;
        if (uncle?.isRed()) {
          grandParent.color = COLORS_TXT.RED;
          uncle.color = COLORS_TXT.BLACK;
          parent.color = COLORS_TXT.BLACK;

          currentNode = grandParent;
        } else {
          if (parent.left === currentNode) {
            currentNode = parent;
            parent = this.rotateLL(parent);
          }
          currentNode = this.rotateRR(grandParent);
          currentNode.color = COLORS_TXT.BLACK;
          grandParent.color = COLORS_TXT.RED;
        }
      }
    }

    this.root && (this.root.color = COLORS_TXT.BLACK);
  }
}
